{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Cloud Detection with OmniCloudMask\n",
    "\n",
    "[![image](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/opengeos/geoai/blob/main/docs/examples/cloud_detection.ipynb)\n",
    "\n",
    "This notebook demonstrates how to use OmniCloudMask integration in GeoAI for detecting clouds and cloud shadows in satellite imagery. OmniCloudMask performs semantic segmentation to classify pixels into four categories:\n",
    "\n",
    "- **0: Clear** - Cloud-free pixels\n",
    "- **1: Thick Cloud** - Opaque cloud cover\n",
    "- **2: Thin Cloud** - Semi-transparent cloud cover\n",
    "- **3: Cloud Shadow** - Shadows cast by clouds\n",
    "\n",
    "OmniCloudMask supports Sentinel-2, Landsat 8, PlanetScope, and Maxar imagery at 10-50m spatial resolution.\n",
    "\n",
    "## Installation\n",
    "\n",
    "Uncomment the following line to install the required packages if needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %pip install -U \"geoai-py[extra]\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Import Libraries"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.colors import ListedColormap\n",
    "import rasterio\n",
    "from rasterio.transform import from_bounds\n",
    "import tempfile\n",
    "import os\n",
    "\n",
    "# Import GeoAI cloud mask utilities\n",
    "# Import from the tools subpackage\n",
    "from geoai.tools.cloudmask import (\n",
    "    predict_cloud_mask,\n",
    "    predict_cloud_mask_from_raster,\n",
    "    predict_cloud_mask_batch,\n",
    "    calculate_cloud_statistics,\n",
    "    create_cloud_free_mask,\n",
    "    CLEAR,\n",
    "    THICK_CLOUD,\n",
    "    THIN_CLOUD,\n",
    "    CLOUD_SHADOW,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Create Synthetic Satellite Imagery\n",
    "\n",
    "For demonstration purposes, let's create synthetic satellite imagery with RGB and NIR bands."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_synthetic_satellite_image(size=(512, 512), cloud_coverage=0.3):\n",
    "    \"\"\"\n",
    "    Create synthetic satellite imagery (R, G, NIR bands).\n",
    "\n",
    "    Args:\n",
    "        size: Image dimensions (height, width)\n",
    "        cloud_coverage: Fraction of image covered by clouds (0-1)\n",
    "\n",
    "    Returns:\n",
    "        3D array with shape (3, height, width) containing R, G, NIR bands\n",
    "    \"\"\"\n",
    "    np.random.seed(42)\n",
    "\n",
    "    # Create base reflectance values typical of vegetation\n",
    "    # Red: low (absorbed by chlorophyll)\n",
    "    # Green: medium\n",
    "    # NIR: high (reflected by vegetation)\n",
    "    red = np.random.rand(*size) * 2000 + 500  # 500-2500\n",
    "    green = np.random.rand(*size) * 3000 + 1000  # 1000-4000\n",
    "    nir = np.random.rand(*size) * 5000 + 3000  # 3000-8000\n",
    "\n",
    "    # Add some spatial structure (vegetation patches)\n",
    "    from scipy.ndimage import gaussian_filter\n",
    "\n",
    "    red = gaussian_filter(red, sigma=20)\n",
    "    green = gaussian_filter(green, sigma=20)\n",
    "    nir = gaussian_filter(nir, sigma=20)\n",
    "\n",
    "    # Add cloud patterns\n",
    "    if cloud_coverage > 0:\n",
    "        # Create cloud mask\n",
    "        cloud_base = np.random.rand(*size)\n",
    "        cloud_base = gaussian_filter(cloud_base, sigma=30)\n",
    "        cloud_mask = cloud_base > (1 - cloud_coverage)\n",
    "\n",
    "        # Clouds have high reflectance in all bands\n",
    "        cloud_value = 8000\n",
    "        red[cloud_mask] = cloud_value + np.random.rand(cloud_mask.sum()) * 1000\n",
    "        green[cloud_mask] = cloud_value + np.random.rand(cloud_mask.sum()) * 1000\n",
    "        nir[cloud_mask] = cloud_value + np.random.rand(cloud_mask.sum()) * 1000\n",
    "\n",
    "    # Stack into (3, H, W) format\n",
    "    image = np.stack([red, green, nir], axis=0).astype(np.float32)\n",
    "\n",
    "    return image\n",
    "\n",
    "\n",
    "# Create synthetic image\n",
    "image = create_synthetic_satellite_image(size=(512, 512), cloud_coverage=0.2)\n",
    "print(f\"Created synthetic image with shape: {image.shape}\")\n",
    "print(f\"Value range: {image.min():.0f} - {image.max():.0f}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Visualize the Input Image\n",
    "\n",
    "Let's visualize the RGB composite of our synthetic satellite image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def visualize_rgb(image, title=\"RGB Composite\"):\n",
    "    \"\"\"\n",
    "    Visualize RGB composite from satellite image.\n",
    "\n",
    "    Args:\n",
    "        image: Array with shape (3, H, W) or (H, W, 3)\n",
    "        title: Plot title\n",
    "    \"\"\"\n",
    "    # Convert to (H, W, 3) if needed\n",
    "    if image.shape[0] == 3:\n",
    "        rgb = image[:3].transpose(1, 2, 0)  # Take R, G, (NIR->B for vis)\n",
    "    else:\n",
    "        rgb = image[:, :, :3]\n",
    "\n",
    "    # Normalize to 0-1 for display\n",
    "    rgb_norm = (rgb - rgb.min()) / (rgb.max() - rgb.min())\n",
    "\n",
    "    plt.figure(figsize=(10, 10))\n",
    "    plt.imshow(rgb_norm)\n",
    "    plt.title(title, fontsize=16)\n",
    "    plt.axis(\"off\")\n",
    "    plt.tight_layout()\n",
    "    plt.show()\n",
    "\n",
    "\n",
    "visualize_rgb(image, \"Synthetic Satellite Image (RGB)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Predict Cloud Mask\n",
    "\n",
    "Now let's use OmniCloudMask to detect clouds and cloud shadows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Predict cloud mask\n",
    "# Note: First run will download the model (may take a moment)\n",
    "cloud_mask = predict_cloud_mask(\n",
    "    image,\n",
    "    batch_size=1,\n",
    "    inference_device=\"cpu\",  # Use \"cuda\" if GPU available\n",
    "    inference_dtype=\"fp32\",  # Use \"bf16\" for faster inference on supported hardware\n",
    "    patch_size=1000,\n",
    "    model_version=3,  # Model versions: 1, 2, or 3\n",
    ")\n",
    "\n",
    "print(f\"Cloud mask shape: {cloud_mask.shape}\")\n",
    "print(f\"Classes found: {np.unique(cloud_mask)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Visualize Cloud Mask\n",
    "\n",
    "Let's visualize the cloud detection results with a color-coded map."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create colormap for cloud classes\n",
    "# Clear (blue), Thick Cloud (white), Thin Cloud (light gray), Shadow (dark gray)\n",
    "colors = [\"#4575b4\", \"#ffffff\", \"#d3d3d3\", \"#404040\"]\n",
    "cmap = ListedColormap(colors)\n",
    "\n",
    "plt.figure(figsize=(12, 10))\n",
    "im = plt.imshow(cloud_mask, cmap=cmap, interpolation=\"nearest\", vmin=0, vmax=3)\n",
    "plt.title(\"Cloud Mask Classification\", fontsize=16)\n",
    "cbar = plt.colorbar(im, ticks=[0, 1, 2, 3])\n",
    "cbar.ax.set_yticklabels([\"Clear\", \"Thick Cloud\", \"Thin Cloud\", \"Shadow\"])\n",
    "plt.xlabel(\"X\")\n",
    "plt.ylabel(\"Y\")\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Calculate Cloud Statistics\n",
    "\n",
    "Let's quantify the cloud coverage and other statistics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stats = calculate_cloud_statistics(cloud_mask)\n",
    "\n",
    "print(\"Cloud Coverage Statistics:\")\n",
    "print(f\"  Total pixels: {stats['total_pixels']:,}\")\n",
    "print(f\"  Clear pixels: {stats['clear_pixels']:,} ({stats['clear_percent']:.1f}%)\")\n",
    "print(f\"  Thick cloud: {stats['thick_cloud_pixels']:,}\")\n",
    "print(f\"  Thin cloud: {stats['thin_cloud_pixels']:,}\")\n",
    "print(f\"  Cloud shadow: {stats['shadow_pixels']:,}\")\n",
    "print(f\"  \\nTotal cloud coverage: {stats['cloud_percent']:.1f}%\")\n",
    "print(f\"  Shadow coverage: {stats['shadow_percent']:.1f}%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Create Cloud-Free Mask\n",
    "\n",
    "Create a binary mask showing which pixels are usable (cloud-free)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Strict cloud-free mask (only clear pixels)\n",
    "cloud_free_strict = create_cloud_free_mask(\n",
    "    cloud_mask, include_thin_clouds=False, include_shadows=False\n",
    ")\n",
    "\n",
    "# Relaxed cloud-free mask (accept thin clouds and shadows)\n",
    "cloud_free_relaxed = create_cloud_free_mask(\n",
    "    cloud_mask, include_thin_clouds=True, include_shadows=True\n",
    ")\n",
    "\n",
    "fig, axes = plt.subplots(1, 2, figsize=(16, 7))\n",
    "\n",
    "axes[0].imshow(cloud_free_strict, cmap=\"RdYlGn\", interpolation=\"nearest\")\n",
    "axes[0].set_title(\n",
    "    f\"Strict Cloud-Free Mask\\n({cloud_free_strict.sum() / cloud_free_strict.size * 100:.1f}% usable)\",\n",
    "    fontsize=14,\n",
    ")\n",
    "axes[0].axis(\"off\")\n",
    "\n",
    "axes[1].imshow(cloud_free_relaxed, cmap=\"RdYlGn\", interpolation=\"nearest\")\n",
    "axes[1].set_title(\n",
    "    f\"Relaxed Cloud-Free Mask\\n({cloud_free_relaxed.sum() / cloud_free_relaxed.size * 100:.1f}% usable)\",\n",
    "    fontsize=14,\n",
    ")\n",
    "axes[1].axis(\"off\")\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Compare Different Classes\n",
    "\n",
    "Let's visualize each cloud class separately."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(2, 2, figsize=(14, 14))\n",
    "axes = axes.flatten()\n",
    "\n",
    "class_names = [\"Clear\", \"Thick Cloud\", \"Thin Cloud\", \"Cloud Shadow\"]\n",
    "class_values = [CLEAR, THICK_CLOUD, THIN_CLOUD, CLOUD_SHADOW]\n",
    "\n",
    "for i, (name, value) in enumerate(zip(class_names, class_values)):\n",
    "    # Create binary mask for this class\n",
    "    class_mask = (cloud_mask == value).astype(np.uint8)\n",
    "    count = class_mask.sum()\n",
    "    percent = count / class_mask.size * 100\n",
    "\n",
    "    axes[i].imshow(class_mask, cmap=\"gray\", interpolation=\"nearest\")\n",
    "    axes[i].set_title(f\"{name}\\n{count:,} pixels ({percent:.1f}%)\", fontsize=12)\n",
    "    axes[i].axis(\"off\")\n",
    "\n",
    "plt.tight_layout()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Working with GeoTIFF Files\n",
    "\n",
    "OmniCloudMask can process GeoTIFF files directly while preserving geospatial metadata."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create temporary directory\n",
    "tmpdir = tempfile.mkdtemp()\n",
    "print(f\"Working directory: {tmpdir}\")\n",
    "\n",
    "# Save synthetic image as GeoTIFF\n",
    "input_tif = os.path.join(tmpdir, \"satellite_image.tif\")\n",
    "output_tif = os.path.join(tmpdir, \"cloud_mask.tif\")\n",
    "\n",
    "# Create geographic transform\n",
    "transform = from_bounds(\n",
    "    west=-120.0,\n",
    "    south=35.0,\n",
    "    east=-119.0,\n",
    "    north=36.0,\n",
    "    width=image.shape[2],\n",
    "    height=image.shape[1],\n",
    ")\n",
    "\n",
    "# Write image to GeoTIFF (3 bands: R, G, NIR)\n",
    "with rasterio.open(\n",
    "    input_tif,\n",
    "    \"w\",\n",
    "    driver=\"GTiff\",\n",
    "    height=image.shape[1],\n",
    "    width=image.shape[2],\n",
    "    count=3,\n",
    "    dtype=image.dtype,\n",
    "    crs=\"EPSG:4326\",\n",
    "    transform=transform,\n",
    "    compress=\"lzw\",\n",
    ") as dst:\n",
    "    for i in range(3):\n",
    "        dst.write(image[i], i + 1)\n",
    "\n",
    "print(f\"Saved satellite image to: {input_tif}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Predict cloud mask from GeoTIFF\n",
    "predict_cloud_mask_from_raster(\n",
    "    input_path=input_tif,\n",
    "    output_path=output_tif,\n",
    "    red_band=1,  # Band indices for R, G, NIR\n",
    "    green_band=2,\n",
    "    nir_band=3,\n",
    "    inference_device=\"cpu\",\n",
    ")\n",
    "\n",
    "print(f\"Cloud mask saved to: {output_tif}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Verify output and metadata preservation\n",
    "with rasterio.open(input_tif) as src_in:\n",
    "    print(\"Input metadata:\")\n",
    "    print(f\"  CRS: {src_in.crs}\")\n",
    "    print(f\"  Transform: {src_in.transform}\")\n",
    "    print(f\"  Bounds: {src_in.bounds}\")\n",
    "    print(f\"  Bands: {src_in.count}\")\n",
    "\n",
    "print()\n",
    "\n",
    "with rasterio.open(output_tif) as src_out:\n",
    "    print(\"Output metadata:\")\n",
    "    print(f\"  CRS: {src_out.crs}\")\n",
    "    print(f\"  Transform: {src_out.transform}\")\n",
    "    print(f\"  Bounds: {src_out.bounds}\")\n",
    "    print(f\"  Bands: {src_out.count}\")\n",
    "\n",
    "    # Read and verify cloud mask\n",
    "    mask_from_file = src_out.read(1)\n",
    "    print(f\"  Classes: {np.unique(mask_from_file)}\")\n",
    "\n",
    "print(\"\\n✓ Geospatial metadata preserved!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 9. Batch Processing Multiple Images\n",
    "\n",
    "Process multiple satellite images at once using batch processing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create multiple test images with different cloud coverage\n",
    "input_files = []\n",
    "cloud_coverages = [0.1, 0.3, 0.5]\n",
    "\n",
    "for i, coverage in enumerate(cloud_coverages):\n",
    "    # Create image with specific cloud coverage\n",
    "    test_image = create_synthetic_satellite_image(\n",
    "        size=(256, 256), cloud_coverage=coverage\n",
    "    )\n",
    "\n",
    "    # Save to file\n",
    "    filepath = os.path.join(tmpdir, f\"scene_{i}_clouds{int(coverage*100)}.tif\")\n",
    "\n",
    "    transform = from_bounds(-120, 35, -119, 36, 256, 256)\n",
    "    with rasterio.open(\n",
    "        filepath,\n",
    "        \"w\",\n",
    "        driver=\"GTiff\",\n",
    "        height=256,\n",
    "        width=256,\n",
    "        count=3,\n",
    "        dtype=test_image.dtype,\n",
    "        crs=\"EPSG:4326\",\n",
    "        transform=transform,\n",
    "    ) as dst:\n",
    "        for j in range(3):\n",
    "            dst.write(test_image[j], j + 1)\n",
    "\n",
    "    input_files.append(filepath)\n",
    "\n",
    "print(f\"Created {len(input_files)} test images\")\n",
    "for f in input_files:\n",
    "    print(f\"  - {os.path.basename(f)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Batch process all images\n",
    "output_dir = os.path.join(tmpdir, \"cloud_masks\")\n",
    "\n",
    "output_files = predict_cloud_mask_batch(\n",
    "    input_paths=input_files,\n",
    "    output_dir=output_dir,\n",
    "    red_band=1,\n",
    "    green_band=2,\n",
    "    nir_band=3,\n",
    "    inference_device=\"cpu\",\n",
    "    suffix=\"_cloudmask\",\n",
    "    verbose=True,\n",
    ")\n",
    "\n",
    "print(f\"\\nProcessed {len(output_files)} images\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Analyze results from batch processing\n",
    "print(\"\\nCloud Coverage Analysis:\")\n",
    "print(\"-\" * 60)\n",
    "\n",
    "for output_file in output_files:\n",
    "    with rasterio.open(output_file) as src:\n",
    "        mask = src.read(1)\n",
    "\n",
    "    stats = calculate_cloud_statistics(mask)\n",
    "\n",
    "    filename = os.path.basename(output_file)\n",
    "    print(f\"\\n{filename}:\")\n",
    "    print(f\"  Clear: {stats['clear_percent']:.1f}%\")\n",
    "    print(f\"  Cloud: {stats['cloud_percent']:.1f}%\")\n",
    "    print(f\"  Shadow: {stats['shadow_percent']:.1f}%\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 10. Use Case: Filter Usable Scenes\n",
    "\n",
    "A common workflow is to filter satellite scenes based on cloud coverage threshold."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 11. Best Practices and Tips\n",
    "\n",
    "### Input Requirements\n",
    "\n",
    "- **Bands**: Requires Red, Green, and NIR bands\n",
    "- **Resolution**: Optimized for 10-50m spatial resolution\n",
    "- **Sensors**: Validated on Sentinel-2, Landsat 8, PlanetScope, Maxar\n",
    "- **Values**: Works with reflectance (0-1) or digital numbers\n",
    "\n",
    "### Performance Tips\n",
    "\n",
    "- Use **inference_dtype='bf16'** for 2-3x speedup on supported hardware\n",
    "- Use **inference_device='cuda'** if GPU available\n",
    "- Adjust **patch_size** based on available memory\n",
    "- Use **batch_size > 1** for faster batch processing\n",
    "\n",
    "### Model Versions\n",
    "\n",
    "- **v3.0** (default): Expanded training dataset for higher accuracy.\n",
    "- **v2.0**: A smaller faster model ensemble with improved robustness.\n",
    "- **v1.0**: Baseline model release supporting the OmniCloudMask paper. \n",
    "\n",
    "### When to Use OmniCloudMask\n",
    "\n",
    "✓ Pre-processing satellite imagery for analysis  \n",
    "✓ Filtering scenes based on cloud coverage  \n",
    "✓ Creating cloud-free composites  \n",
    "✓ Quality assessment of satellite data  \n",
    "✓ Time series analysis (exclude cloudy observations)  \n",
    "\n",
    "### Sensor-Specific Band Indices\n",
    "\n",
    "**Sentinel-2:**\n",
    "```python\n",
    "red_band=4, green_band=3, nir_band=8\n",
    "```\n",
    "\n",
    "**Landsat 8/9:**\n",
    "```python\n",
    "red_band=4, green_band=3, nir_band=5\n",
    "```\n",
    "\n",
    "**PlanetScope:**\n",
    "```python\n",
    "red_band=3, green_band=2, nir_band=4\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "In this notebook, we demonstrated:\n",
    "\n",
    "1. ✅ Predicting cloud masks from numpy arrays\n",
    "2. ✅ Visualizing cloud detection results\n",
    "3. ✅ Calculating cloud coverage statistics\n",
    "4. ✅ Creating cloud-free masks\n",
    "5. ✅ Processing GeoTIFF files while preserving metadata\n",
    "6. ✅ Batch processing multiple scenes\n",
    "7. ✅ Filtering scenes based on cloud coverage\n",
    "\n",
    "OmniCloudMask is a powerful tool for cloud detection in satellite imagery, essential for many remote sensing workflows.\n",
    "\n",
    "## References\n",
    "\n",
    "- [OmniCloudMask GitHub Repository](https://github.com/DPIRD-DMA/OmniCloudMask)\n",
    "- [GeoAI Documentation](https://opengeoai.org)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Cleanup temporary files\n",
    "import shutil\n",
    "\n",
    "shutil.rmtree(tmpdir)\n",
    "print(\"Cleaned up temporary files\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "geo",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
